# utils.py - 工具函数模块
import os
import time
import numpy as np
import torch
import dill
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
import matplotlib.pyplot as plt
import seaborn as sns
import csv


# 设置随机种子
def same_seeds(seed):
    """
    设置所有随机种子以确保实验可重复性
    
    Args:
        seed (int): 随机种子值
        
    功能说明:
        - 设置PyTorch、NumPy等库的随机种子
        - 确保每次运行实验时得到相同的结果
        - 启用CUDA确定性算法以增强GPU上的可重复性
    """
    # 设置PyTorch的随机种子
    torch.manual_seed(seed)
    # 如果CUDA可用，设置CUDA相关的随机种子
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)      # 为当前GPU设置随机种子
        torch.cuda.manual_seed_all(seed)  # 为所有GPU设置随机种子
    # 设置NumPy的随机种子
    np.random.seed(seed)
    # 启用cudnn基准模式，可以加速训练但可能略微降低可重复性
    torch.backends.cudnn.benchmark = True


# 早停机制
class EarlyStopping:
    """
    早停机制类，用于防止过拟合
    
    属性:
        save_path (str): 模型保存路径
        patience (int): 容忍的验证损失不改善的轮数
        verbose (bool): 是否打印详细信息
        counter (int): 当前计数器，记录验证损失不改善的轮数
        best_score (float): 最佳验证分数（负的验证损失）
        early_stop (bool): 是否应该早停
        val_loss_min (float): 最小验证损失
        delta (float): 改善的最小阈值，只有超过此阈值才认为是改善
    """
    def __init__(self, save_path, patience=7, verbose=False, delta=0):
        """
        初始化早停机制
        
        Args:
            save_path (str): 模型保存路径
            patience (int): 容忍的验证损失不改善的轮数，默认为7
            verbose (bool): 是否打印详细信息，默认为False
            delta (float): 改善的最小阈值，默认为0
        """
        self.save_path = save_path
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.inf  # 初始化为无穷大
        self.delta = delta

    def __call__(self, val_loss, model):
        """
        调用早停机制，根据验证损失决定是否保存模型或停止训练
        
        Args:
            val_loss (float): 当前验证损失
            model: 要保存的模型
        """
        # 将验证损失转换为分数（负值，因为我们要最小化损失）
        score = -val_loss
        
        # 如果是第一次调用，保存当前模型
        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
        # 如果当前分数没有比最佳分数改善（考虑阈值）
        elif score < self.best_score + self.delta:
            self.counter += 1  # 增加计数器
            print(f"EarlyStopping counter: {self.counter} out of {self.patience}")
            # 如果计数器超过容忍值，触发早停
            if self.counter >= self.patience:
                self.early_stop = True
        # 如果当前分数比最佳分数改善（考虑阈值）
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
            self.counter = 0  # 重置计数器

    def save_checkpoint(self, val_loss, model):
        """
        保存模型检查点
        
        Args:
            val_loss (float): 验证损失
            model: 要保存的模型
        """
        # 如果verbose为True，打印验证损失改善信息
        if self.verbose:
            print(f"Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...")
        # 获取保存路径
        path = self.save_path
        # 使用dill保存模型（dill可以保存更复杂的Python对象）
        torch.save(model, path, pickle_module=dill)
        # 更新最小验证损失
        self.val_loss_min = val_loss


# 评估指标
def metrics_of_pv(preds, trues):
    """
    计算预测值与真实值之间的各种评估指标
    
    Args:
        preds: 预测值数组
        trues: 真实值数组
        
    Returns:
        list: 包含四个评估指标的列表 [MAE, MSE, R2, MBE]
            - MAE: 平均绝对误差
            - MSE: 均方误差
            - R2: 决定系数
            - MBE: 平均偏差误差
    """
    # 将预测值和真实值转换为NumPy数组并展平
    pred = np.array(preds).flatten()
    true = np.array(trues).flatten()
    # 计算平均绝对误差（Mean Absolute Error）
    mae = np.round(mean_absolute_error(true, pred), 4)
    # 计算均方误差（Mean Squared Error）
    mse = np.round(mean_squared_error(true, pred), 4)
    # 计算决定系数（R-squared）
    r2 = np.round(r2_score(true, pred), 4)
    # 计算平均偏差误差（Mean Bias Error）
    mbe = np.round(np.mean(pred - true), 4)
    return [mae, mse, r2, mbe]


# 可视化函数
def plot_training_loss(train_losses, valid_losses, model_name, dataset):
    """
    绘制训练损失和验证损失曲线
    
    Args:
        train_losses (list): 训练损失列表
        valid_losses (list): 验证损失列表
        model_name (str): 模型名称，用于保存文件
        dataset (str): 数据集名称，用于保存文件
    """
    # 创建图形，设置大小
    plt.figure(figsize=(10, 6))
    # 绘制训练损失曲线
    plt.plot(np.arange(len(train_losses)), train_losses, label="train loss")
    # 绘制验证损失曲线
    plt.plot(np.arange(len(valid_losses)), valid_losses, label="valid loss")
    # 添加图例
    plt.legend()
    # 添加x轴标签
    plt.xlabel("epochs")
    # 添加标题
    plt.title("Train_loss & Valid_loss")
    # 保存图像
    plt.savefig(f"data_record/{dataset}/{model_name}_loss.png")
    # 显示图像
    plt.show()


def plot_time_series_prediction(test_labels, test_preds, timestamp_test, time_length, model_name, dataset):
    """
    绘制时间序列预测对比图
    
    Args:
        test_labels: 真实值数组
        test_preds: 预测值数组
        timestamp_test: 测试集时间戳
        time_length: 输入序列长度
        model_name (str): 模型名称，用于保存文件
        dataset (str): 数据集名称，用于保存文件
    """
    # 创建图形，设置大小
    plt.figure(figsize=(12, 6))
    # 从时间戳中提取日期值，跳过前time_length个点（因为输入序列长度）
    time_axis = timestamp_test['date'].values[time_length:]
    # 截取与测试标签长度相同的时间轴
    time_axis = time_axis[:len(test_labels)]
    # 绘制真实值曲线
    plt.plot(time_axis, test_labels[:, 0], label="True Values", color='blue')
    # 绘制预测值曲线（虚线）
    plt.plot(time_axis, test_preds[:, 0], label="Predicted Values", color='orange', linestyle='--')
    # 添加图例
    plt.legend()
    # 添加x轴标签
    plt.xlabel("Time")
    # 添加y轴标签
    plt.ylabel("Values")
    # 添加标题
    plt.title("Time Series: True vs Predicted Values")
    # 旋转x轴标签45度，防止重叠
    plt.xticks(rotation=45)
    # 自动调整布局，防止标签被截断
    plt.tight_layout()
    # 保存图像
    plt.savefig(f"data_record/{dataset}/{model_name}_true_pred.png")
    # 显示图像
    plt.show()


def plot_error_distribution(test_labels, test_preds, model_name, dataset):
    """
    绘制预测误差分布图
    
    Args:
        test_labels: 真实值数组
        test_preds: 预测值数组
        model_name (str): 模型名称，用于保存文件
        dataset (str): 数据集名称，用于保存文件
    """
    # 创建图形，设置大小
    plt.figure(figsize=(10, 6))
    # 计算预测误差（预测值 - 真实值）
    errors = test_preds[:, 0] - test_labels[:, 0]
    # 绘制误差分布直方图，并添加核密度估计曲线
    sns.histplot(errors, kde=True, color='purple')
    # 添加零误差参考线（红色虚线）
    plt.axvline(x=0, color='red', linestyle='--', label='Zero Error')
    # 添加平均误差参考线（绿色实线）
    plt.axvline(x=np.mean(errors), color='green', linestyle='-', label='Mean Error')
    # 添加图例
    plt.legend()
    # 添加x轴标签
    plt.xlabel("Prediction Error")
    # 添加标题
    plt.title("Error Distribution")
    # 保存图像
    plt.savefig(f"data_record/{dataset}/{model_name}_error_dist.png")
    # 显示图像
    plt.show()


def plot_cumulative_error(test_labels, test_preds, model_name, dataset):
    """
    绘制累积误差图
    
    Args:
        test_labels: 真实值数组
        test_preds: 预测值数组
        model_name (str): 模型名称，用于保存文件
        dataset (str): 数据集名称，用于保存文件
    """
    # 创建图形，设置大小
    plt.figure(figsize=(10, 6))
    # 计算预测误差（预测值 - 真实值）
    errors = test_preds[:, 0] - test_labels[:, 0]
    # 计算累积误差
    cum_errors = np.cumsum(errors)
    # 绘制累积误差曲线
    plt.plot(cum_errors, color='purple')
    # 添加零参考线（红色虚线）
    plt.axhline(y=0, color='red', linestyle='--')
    # 添加x轴标签
    plt.xlabel("Sample Index")
    # 添加y轴标签
    plt.ylabel("Cumulative Error")
    # 添加标题
    plt.title("Cumulative Error Plot")
    # 保存图像
    plt.savefig(f"data_record/{dataset}/{model_name}_cum_error.png")
    # 显示图像
    plt.show()


def plot_daily_power_curve(test_labels, test_preds, timestamp_test, time_length, model_name, dataset):
    """
    绘制一天功率曲线图
    
    Args:
        test_labels: 真实值数组
        test_preds: 预测值数组
        timestamp_test: 测试集时间戳
        time_length: 输入序列长度
        model_name (str): 模型名称，用于保存文件
        dataset (str): 数据集名称，用于保存文件
    """
    # 创建图形，设置大小
    plt.figure(figsize=(12, 6))
    # 每天的时间点数（24小时）
    points_per_day = 24

    # 获取测试集时间戳，跳过前time_length个点
    time_axis = timestamp_test['date'][time_length:time_length + len(test_labels)]
    
    # 查找一天的开始点（00:00）
    start_idx = None
    for i, dt in enumerate(time_axis):
        if dt.hour == 0 and dt.minute == 0:
            start_idx = i
            break

    # 如果没有找到00:00的时间点，使用测试集第一天数据
    if start_idx is None:
        print("未找到00:00的时间点，使用测试集第一天数据")
        start_idx = 0

    # 计算一天的结束点
    end_idx = start_idx + points_per_day
    # 如果数据不足以覆盖完整一天，使用可用数据
    if end_idx > len(test_labels):
        print(f"数据不足以覆盖完整一天，使用可用数据")
        end_idx = len(test_labels)

    # 提取一天的真实值和预测值
    day_labels = test_labels[start_idx:end_idx, 0]
    day_preds = test_preds[start_idx:end_idx, 0]
    # 提取一天的时间轴
    day_time_axis = time_axis[start_idx:end_idx]
    # 将时间转换为"HH:MM"格式
    hours = [dt.strftime('%H:%M') for dt in day_time_axis]

    # 如果时间点不足24小时，打印警告
    if len(hours) < points_per_day:
        print(f"时间点不足24小时，仅有{len(hours)}个点")

    # 绘制真实值曲线
    plt.plot(hours, day_labels, label="True Values", color='blue')
    # 绘制预测值曲线（虚线）
    plt.plot(hours, day_preds, label="Predicted Values", color='orange', linestyle='--')
    # 添加图例
    plt.legend()
    # 添加x轴标签
    plt.xlabel("Time of Day")
    # 添加y轴标签
    plt.ylabel("Power")
    # 添加标题
    plt.title("Daily Power Curve: True vs Predicted")
    # 设置x轴刻度，每个时间点都显示，旋转45度防止重叠
    plt.xticks(ticks=range(0, len(hours), 1), labels=[hours[i] for i in range(0, len(hours), 1)], rotation=45, fontsize=8)
    # 自动调整布局，防止标签被截断
    plt.tight_layout()
    # 保存图像
    plt.savefig(f"data_record/{dataset}/{model_name}_daily_power_curve.png")
    # 显示图像
    plt.show()


def save_metrics_to_csv(site, model_name, metrics, dataset):
    """
    将评估指标保存到CSV文件
    
    Args:
        site (str): 站点名称
        model_name (str): 模型名称
        metrics (list): 评估指标列表 [MAE, MSE, R2, MBE]
        dataset (str): 数据集名称，用于确定保存路径
    """
    # 打开CSV文件，使用追加模式
    with open(f'data_record/{dataset}/Metrics_{model_name}.csv', 'a', encoding='utf-8', newline='') as f:
        # 创建CSV写入器
        csv_write = csv.writer(f)
        # 写入一行数据：站点名称和四个评估指标
        csv_write.writerow([f'{site}_pred1_{model_name}', metrics[0], metrics[1], metrics[2], metrics[3]])